'''
Created on Sep 23, 2011

@author: oabalbin
'''
import os
import sys
import numpy as np
import exome.vcftools.SNP_positions as snppos
from optparse import OptionParser
from collections import defaultdict,deque

def isec_of_snvs(snvs_by_tool, snvs_annot_dbs):
    '''
    Returns the intersection between snvs in both of the dictionaries
    '''
    pm = defaultdict()
    
    for db, snv_db in snvs_annot_dbs.iteritems():
        for tool, snvs in snvs_by_tool.iteritems():
            pm[tool+'&'+db] = set(snvs).intersection(set(snv_db))
            
    return pm

def total_dict_creator(list_of_dicts):
    '''
    '''
    total_dict = defaultdict()
    for d in list_of_dicts:
        for k, snvs in d.iteritems():
            total_dict[k]=len(snvs)
    return total_dict
            
    
def main_concesus(vcf_files, annot_files, ofile, mmapfile='',complement=False):
    '''
    '''
    snvs_union = set()
    pfiles=defaultdict()
    snvs_dict = defaultdict()
    snvs_annot_dbs = defaultdict()
    all_dicts=deque()
    
    for ft, f in vcf_files.iteritems():
        # Generate a positions files for query vcf.
        print f
        fpos = f+'.pos'
        if ft =='VARS':
            snppos.make_posbase_file_varsc(f,fpos)
        else:
            snppos.make_posbase_file(f,fpos)
        snvs_dict[ft]=snppos.read_positions(fpos+'.base')
        
    for af, afile in annot_files.iteritems():
        snvs_annot_dbs[af] = snppos.read_positions(afile)   
    
    # Intersection of snvs and annotated snvs
    #annotated_snvs = isec_snvs_annot(snvs_dict, snvs_annot_dbs)
    annotated_snvs = isec_of_snvs(snvs_dict, snvs_annot_dbs)
    tools_isec = isec_of_snvs(snvs_dict, snvs_dict)
    # Report table:
    all_dicts.append(snvs_dict),all_dicts.append(snvs_annot_dbs)
    all_dicts.append(annotated_snvs), all_dicts.append(tools_isec)
    total_dict=total_dict_creator(all_dicts)

    return total_dict
    
    


if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options] ")
    '''
    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)
    # Specific options
    optionparser.add_option("-f", "--vcf_files", action="append", type="str",
                            dest="vcf_files",
                            help="others_vcf_files. Pass vcf file in the following order: GATK, SAM, Other")
    '''
    optionparser.add_option("-a", "--annot_files", action="append", type="str",
                            dest="annot_files",
                            help="other annotation files. ")
    optionparser.add_option("-o", "--output_file", dest="output_file",
                            help="output_file")
    optionparser.add_option("-c", "--complement", action='store_true', default=False,
                            help="Not used")    
    optionparser.add_option("-d", "--dirpath", dest='dirpath',
                            help="root directory")
    optionparser.add_option("-N", "--N_tools", dest='N_tools',type="int",
                            help="number of tools to compare")
    

    (options, args) = optionparser.parse_args()
    outfile =options.output_file+'.txt'
    logfile =open(options.output_file+'.log','w')
    mmapfile =options.output_file+'.memmap'
    print options
    '''
    if len(options.vcf_files) < 2:
        print >> sys.stderr, "You need a minimum of 2 vcf files to do the analysis"
        sys.exit(0)
    elif len(options.vcf_files) ==3:
        vcf_files = {'GATK':options.vcf_files[0],'SAM':options.vcf_files[1], 
                     'VARS':options.vcf_files[2]}
    elif len(options.vcf_files) ==2:
        vcf_files = {'GATK':options.vcf_files[0],'SAM':options.vcf_files[1]}
    '''
    
    '''
    /exds/projects/alignment_indexes/gatk/hg19/dbSNP/dbSNP132_12_03_10_00-All_processed_04_05_11_2_db132.pos.base
    /exds/projects/alignment_indexes/gatk/hg19/hapmap3.3.genotypes.b37_fwd_process_tab2_NA12878_b37.pos.base
    '''
    annot_files_dict={'dbSNP':options.annot_files[0],'HapMap':options.annot_files[1]}
    N_tools=options.N_tools
    # Do the work
    print options.dirpath
    dirList = os.listdir(options.dirpath)
    nrows = N_tools*N_tools + N_tools*len(annot_files_dict) + len(annot_files_dict) + N_tools
    ncolumns  = len(dirList)
    master_table=np.zeros((nrows,ncolumns))
    sample_list=deque()
    for j, d in enumerate(dirList):
        
        dir=os.path.join(options.dirpath,d)
        anadir=os.path.join(dir,'analysis')
        if os.path.isdir(dir) and os.path.exists(anadir):
            gatk_file=os.path.join(anadir,'gatk_calls/'+d+'.gatk_snps.b.filtered.model.vcf')
            sam_file=os.path.join(anadir,'sam_calls_quick/'+d+'benign.sam_snps.mpileup.vcf')
            vars_file=os.path.join(anadir,'varscan_calls/'+d+'_snps_varscan_single')
            
            sample_list.append(str(j)+'_'+d)
            if not os.path.exists(gatk_file) or not os.path.exists(sam_file) or not os.path.exists(vars_file):
                logfile.write(",".join([str(os.path.exists(gatk_file)), gatk_file,'\n']))
                logfile.write(",".join([str(os.path.exists(sam_file)), sam_file,'\n']))
                logfile.write(",".join([str(os.path.exists(vars_file)), vars_file,'\n']))
                continue
            
            
            vcf_files = {'GATK':gatk_file,'SAM':sam_file,'VARS':vars_file}
            #print d, vcf_files
            
            total_dict=main_concesus(vcf_files, annot_files_dict, outfile, mmapfile, options.complement)
            skeys=total_dict.keys()
            skeys.sort()
            for i,s in enumerate(skeys):
                master_table[i,j] = total_dict[s]
    
    labels_list=deque()
    for i,s in enumerate(skeys):
        labels_list.append(str(i)+'_'+s)
    
    outfile=open(outfile,'w')
    outfile.write("Comparison"+'\t'+",".join(sample_list).replace(',','\t')+'\n')
    for i in range(master_table.shape[0]):
        outfile.write(labels_list[i]+'\t'+",".join(map(str,list(master_table[i,:]))).replace(",","\t")+'\n')
        
    outfile.close()
    logfile.close()
